package com.dragon.datastructure.tree.avl;

import java.util.ArrayList;

/**
 * AVLTree
 * 是平衡二叉树，而平衡二叉树的定义为: 对于每一个节点，都有 这个节点的 左子树 与 右子树的高度差的绝对值不超过1
 * <p>
 * 二叉树中的高度的定义: 是从下往上数，如果一直到根节点，代表是整个树的高度，可以把叶子节点的高度定义为1，也可以定义为0，而根节点的高度代表整棵树的高度
 * 二叉树中的深度的定义: 是从根节点开始从上往下数，从根节点到叶子节点的最大的路径树代表整棵树的深度，而一颗树的深度与一棵树的高度是相等的
 *
 * @author dragon
 * @since 2020/1/3
 */
public class AvlTree<E extends Comparable<E>> {

    private int size;

    private Node root;

    private class Node {
        /**
         * 节点的元素内容
         */
        E item;
        /**
         * 节点所在的高度
         */
        int height;

        Node left;
        Node right;

        public Node(E item, Node left, Node right) {
            this.item = item;
            this.left = left;
            this.right = right;
        }

        public Node() {
            height = 1;
        }

    }

    /**
     * 查询元素个数
     *
     * @return
     */
    public int size() {

        return size;
    }

    public boolean isEmpty() {
        return size == 0;
    }

    /**
     * 返回node节点的高度, 如果node节点为null，则返回 0
     *
     * @param node
     * @return
     */
    public int height(Node node) {
        if (node == null) {
            return 0;
        }

        return node.height;
    }

    /**
     * 返回一个节点的平衡因子，平衡因子代表的是一个节点的左子树的高度与右子树的高度的差值
     *
     * @param node
     * @return
     */
    public int balanceFactor(Node node) {
        if (node == null) {
            return 0;
        }

        return height(node.left) - height(node.right);
    }

    /**
     * 判断当前的AVL树是否是一颗平衡树
     *
     * @return
     */
    public boolean isBalanced() {

        return isBalanced(root);
    }


    private boolean isBalanced(Node node) {
        if (node == null) {
            return true;
        }

        int balanceFactor = balanceFactor(node);
        if (Math.abs(balanceFactor) > 1) {
            return false;
        }

        return isBalanced(node.left) && isBalanced(node.right);
    }

    /**
     * 判断该树是否是一颗二分搜素树
     *
     * @return
     */
    public boolean isBinarySearchTree() {
        if (root == null) {
            return true;
        }
        ArrayList<Node> list = new ArrayList<>();

        midOrder(root, list);

        for (int i = 1; i < list.size(); i++) {
            if (list.get(i - 1).item.compareTo(list.get(i).item) > 0) {
                return false;
            }
        }

        return true;
    }

    private void midOrder(Node node, ArrayList<Node> keys) {
        if (node == null) {
            return;
        }
        midOrder(node.left, keys);
        keys.add(node);
        midOrder(node.right, keys);
    }

    /**
     * 添加元素
     *
     * @param e
     */
    public void add(E e) {
        root = addRec(root, e);
    }

    private Node addRec(Node root, E e) {
        if (root == null) {
            size++;
            return new Node(e, null, null);
        }
        if (root.item.compareTo(e) > 0) {
            root.left = addRec(root.left, e);
        } else if (root.item.compareTo(e) < 0) {
            root.right = addRec(root.right, e);
        } else {
            root.item = e;
        }
        // 更新height
        root.height = Math.max(height(root.left), height(root.right)) + 1;
        // 计算平衡因子
        int balanceFactor = balanceFactor(root);

        /*
         * LL
         * balanceFactor > 1:代表这个节点不平衡，该节点左子树的高度是大于该节点右子树的高度
         * balanceFactor(node.left) >=0 : 代表新插入的节点在 不平衡节点的左子树的左子树，也就是在不平衡节点的左侧的左侧，记做LL
         */
        if (balanceFactor > 1 && balanceFactor(root.left) >= 0) {
            // 右旋转
            return rightRotate(root);
        }
        /*
         * RR
         * balanceFactor < -1: 代表这个节点不平衡且该节点右子树的高度是大于该节点左子树的高度，说明这棵树是往右倾斜
         * balanceFactor(node.right) >=0 : 代表新插入的节点在 不平衡节点的右子树的右子树，也就是不平衡节点的右侧的右侧，记做RR
         */
        if (balanceFactor < -1 && balanceFactor(root.right) <= 0) {
            // 左旋转
            return leftRotate(root);
        }
        /*
         * LR
         * balanceFactor > 1:代表这个节点不平衡，该节点左子树的高度是大于该节点右子树的高度
         * balanceFactor(node.left) <0 : 代表新插入的节点在 不平衡节点的左子树的右子树，也就是不平衡节点的左侧的右侧，记做LR
         */
        if (balanceFactor > 1 && balanceFactor(root.left) < 0) {
            // 先左旋
            root.left = leftRotate(root.left);
            // 在右旋
            return rightRotate(root);
        }

        //RL
        if (balanceFactor < -1 && balanceFactor(root.right) > 0) {
            // 先右旋
            root.right = rightRotate(root.right);
            // 在左旋
            return leftRotate(root);
        }

        return root;

    }

    public void remove(E e) {
        root = removeRec(root, e);
    }

    private Node removeRec(Node node, E e) {
        if (node == null) {
            return null;
        }
        Node retNode;
        if (node.item.compareTo(e) > 0) {
            node.left = removeRec(node.left, e);
            retNode = node;
        } else if (node.item.compareTo(e) < 0) {
            node.right = removeRec(node.right, e);
            retNode = node;
        } else {
            if (node.left == null) {
                Node rightNode = node.right;
                node.right = null;
                size--;
                retNode = rightNode;
            } else if (node.right == null) {
                Node leftNode = node.left;
                node.left = null;
                size--;
                retNode = leftNode;
            } else {
                Node successor = minimum(node.right);
                successor.right = removeRec(node.right, successor.item);
                successor.left = node.left;
                node.left = node.right = null;
                retNode = successor;
            }
        }

        if (retNode == null) {
            return null;
        }

        // 更新height
        retNode.height = Math.max(height(retNode.left), height(retNode.right)) + 1;
        // 计算平衡因子
        int balanceFactor = balanceFactor(retNode);

        /*
         * LL
         * balanceFactor > 1:代表这个节点不平衡，该节点左子树的高度是大于该节点右子树的高度
         * balanceFactor(node.left) >=0 : 代表新插入的节点在 不平衡节点的左子树的左子树，也就是在不平衡节点的左侧的左侧，记做LL
         */
        if (balanceFactor > 1 && balanceFactor(retNode.left) >= 0) {
            // 右旋转
            return rightRotate(retNode);
        }

        /*
         * RR
         * balanceFactor < -1: 代表这个节点不平衡且该节点右子树的高度是大于该节点左子树的高度，说明这棵树是往右倾斜
         * balanceFactor(node.right) >=0 : 代表新插入的节点在 不平衡节点的右子树的右子树，也就是不平衡节点的右侧的右侧，记做RR
         */
        if (balanceFactor < -1 && balanceFactor(retNode.right) <= 0) {
            // 左旋转
            return leftRotate(retNode);
        }

        /*
         * LR
         * balanceFactor > 1:代表这个节点不平衡，该节点左子树的高度是大于该节点右子树的高度
         * balanceFactor(node.left) <0 : 代表新插入的节点在 不平衡节点的左子树的右子树，也就是不平衡节点的左侧的右侧，记做LR
         */
        if (balanceFactor > 1 && balanceFactor(retNode.left) < 0) {
            // 先左旋
            retNode.left = leftRotate(retNode.left);
            // 在右旋
            return rightRotate(retNode);
        }

        //RL
        if (balanceFactor < -1 && balanceFactor(retNode.right) > 0) {
            // 先右旋
            retNode.right = rightRotate(retNode.right);
            // 在左旋
            return leftRotate(retNode);
        }

        return retNode;
    }

    private Node minimum(Node node) {
        if (node == null) {
            return null;
        }
        Node cur = node;
        while (cur.left != null) {
            cur = cur.left;
        }

        return cur;
    }

    // 对节点y进行向右旋转操作，返回旋转后新的根节点x
    //        y                              x
    //       / \                           /   \
    //      x   T4     向右旋转 (y)        z     y
    //     / \       - - - - - - - ->    / \   / \
    //    z   T3                       T1  T2 T3 T4
    //   / \
    // T1   T2
    private Node rightRotate(Node y) {
        Node x = y.left;
        Node t3 = x.right;
        x.right = y;
        y.left = t3;

        y.height = Math.max(height(y.left), height(y.right)) + 1;
        x.height = Math.max(height(x.left), height(x.right)) + 1;

        return x;
    }

    // 对节点y进行向左旋转操作，返回旋转后新的根节点x
    //    y                             x
    //  /  \                          /   \
    // T1   x      向左旋转 (y)       y     z
    //     / \   - - - - - - - ->   / \   / \
    //   T2  z                     T1 T2 T3 T4
    //      / \
    //     T3 T4
    private Node leftRotate(Node y) {
        Node x = y.right;
        Node t3 = x.left;
        x.left = y;
        y.right = t3;

        y.height = Math.max(height(y.left), height(y.right)) + 1;
        x.height = Math.max(height(x.left), height(x.right)) + 1;

        return x;
    }


}
